package org.ricks.net.action.rpc;

import org.ricks.net.action.context.FlowContext;
import org.ricks.net.action.context.head.RpcHeadMeta;
import org.ricks.net.action.context.session.TcpSessionAdapter;
import org.ricks.net.action.parser.ProtocolParser;
import org.ricks.net.action.parser.TypeReference;
import org.ricks.net.action.utils.RpcException;
import org.ricks.net.timer.HashedWheelTimer;

import java.lang.reflect.Type;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.LongAdder;

// RpcClient.java
public class RpcClient {

    private final HashedWheelTimer timer;
    private final TcpSessionAdapter session;
    private final AtomicInteger requestIdGenerator = new AtomicInteger();
    private final ConcurrentHashMap<Integer, PendingRequest<?>> pendingRequests = new ConcurrentHashMap<>();
    private final LongAdder counter = new LongAdder();
    private final long timeoutMillis;
    private final RpcMetrics metrics = new RpcMetrics();
    private static final int MAX_PENDING_REQUESTS = 200_0000;


    // 优化2：添加原子状态管理
    private static class PendingRequest<T> {
        final int requestId;
        final long startTime;
        final CompletableFuture<T> future;
        final Type resultType;
        final AtomicBoolean completed = new AtomicBoolean(false);

        PendingRequest(int requestId, long startTime, CompletableFuture<T> future,
                       Type resultType) {
            this.requestId = requestId;
            this.startTime = startTime;
            this.future = future;
            this.resultType = resultType;
        }
    }


    // 优化7: 添加序列化线程池
    public RpcClient(TcpSessionAdapter session,  long timeoutMillis) {
        this.session = session;
        this.timeoutMillis = timeoutMillis;

        // 创建时间轮定时器
        this.timer = new HashedWheelTimer(
                r -> new Thread(r, "RpcTimeoutTimer"),
                1, // 1ms精度
                512 // 时间轮大小
        );

    }

    /**
     * 处理RPC响应
     */
    public void handleRpcResponse(FlowContext<RpcHeadMeta> context) {
        RpcHeadMeta headMeta = context.headMeta;
        if (!headMeta.isRpcResponse()) return;

        int requestId = headMeta.rpcRequestId();
        PendingRequest<?> pending = pendingRequests.remove(requestId);
        if (pending == null) {
            metrics.recordDroppedResponse();
            return;
        }

        // 原子标记完成状态
        if (!pending.completed.compareAndSet(false, true)) {
            return; // 已由超时处理
        }

        counter.decrement();
        long latency = System.nanoTime() - pending.startTime;

        try {
            ProtocolParser parser = session.getParser();
            Object result = parser.deserialize(context.getData(), pending.resultType);
            if (result instanceof RpcError error) {
                metrics.recordResponse(latency, false);
                metrics.recordInvokeError();
                pending.future.completeExceptionally(
                        new RpcException(error.getMessage())
                );
            } else {
                metrics.recordResponse(latency, true);
                completeFutureSafely(pending, result);
            }
        } catch (Exception e) {
            metrics.recordDecodeError();
            pending.future.completeExceptionally(e);
        }
    }

    // 修复3: 添加类型安全的完成方法
    @SuppressWarnings("unchecked")
    private <T> void completeFutureSafely(PendingRequest<T> pending, Object result) {
        try {
            T typedResult = (T) result;
            pending.future.complete(typedResult);
        } catch (ClassCastException e) {
            String errorMsg = String.format(
                    "Type mismatch for request %d. Expected %s but received %s",
                    pending.requestId,  // 需在 PendingRequest 添加 requestId
                    pending.resultType.getTypeName(),
                    result.getClass().getName()
            );
            pending.future.completeExceptionally(new ClassCastException(errorMsg));
            metrics.recordTypeError();
        }
    }

    // 优化8: 异步序列化
    public <T> CompletableFuture<T> invoke(short cmd, Object arg, Class<T> resultType) {
        return invoke(cmd, arg, (Type) resultType, timeoutMillis, TimeUnit.MILLISECONDS);
    }

    // 修改invoke方法支持自定义超时
    public <T> CompletableFuture<T> invoke(short cmd, Object arg, Class<T> resultType, long timeout, TimeUnit unit) {
        return invoke(cmd, arg, (Type) resultType, timeout, unit);
    }

    public <T> CompletableFuture<T> invoke(short cmd, Object arg, TypeReference<T> typeRef) {
        return invoke(cmd, arg, typeRef.getType(), timeoutMillis, TimeUnit.MILLISECONDS);
    }

    private <T> CompletableFuture<T> invoke(short cmd, Object arg, Type resultType,  long timeout, TimeUnit unit) {
        // 检查系统负载
        if (counter.sum() > MAX_PENDING_REQUESTS) {
            CompletableFuture<T> failed = new CompletableFuture<>();
            failed.completeExceptionally(new RpcException("Server busy"));
            metrics.recordRejectedRequest();
            return failed;
        }

        int requestId = requestIdGenerator.incrementAndGet();

        CompletableFuture<T> future = new CompletableFuture<>();

        PendingRequest<T> pending = new PendingRequest<>(requestId, System.nanoTime(), future, resultType);

        // 注册超时检测
        timer.schedule(() -> handleTimeout(pending), timeoutMillis, TimeUnit.MILLISECONDS);

        pendingRequests.put(requestId, pending);
        counter.increment();
        metrics.recordTotalRequests();

        // 发送请求
        RpcHeadMeta headMeta = new RpcHeadMeta(cmd, false, false, requestId);
        ProtocolParser parser = session.getParser();
        try {
            byte[] body = parser.serialize(arg);
            FlowContext<RpcHeadMeta> context = new FlowContext<>(session, headMeta, body, parser);
            context.writeBytes(body);
        } catch (Exception e) {
            handleSendError(requestId, pending, e);
        }

        return future;
    }

    /**
     * 处理发送错误
     */
    private <T> void handleSendError(int requestId,
                                     PendingRequest<T> pending, Exception e) {
        if (pendingRequests.remove(requestId, pending)) {
            counter.decrement();
            if (pending.completed.compareAndSet(false, true)) {
                pending.future.completeExceptionally(e);
            }
        }
    }

    /**
     * 处理请求超时
     */
    private <T> void handleTimeout(PendingRequest<T> pending) {
        if (pending.completed.compareAndSet(false, true)) {
            if (pendingRequests.remove(pending.requestId, pending)) {
                counter.decrement();
                metrics.recordTimeout();
                pending.future.completeExceptionally(
                        new TimeoutException("RPC call timed out - requestId:" + pending.requestId)
                );
            }
        }
    }

    /**
     * 关闭客户端
     */
    public void close() {
        timer.shutdown();

        // 清空所有分片
        for (PendingRequest pending : pendingRequests.values()) {
            if (pending.completed.compareAndSet(false, true)) {
                pending.future.completeExceptionally(new RpcException( "RPC client closed"));
            }
        };
        pendingRequests.clear();
        counter.reset();
    }

    public void printStats () {
        metrics.printStats();
        System.out.println("消息累计 ：" + counter.sum());
    }
}